package ru.compft.web;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.IOException;
import java.io.PrintWriter;

/**
 * User: ASUS
 * Date: 24.08.2011
 * Time: 0:01:51
 * Обертка для записи времени генерации ответа в response
 */
public class CountingServletResponse extends HttpServletResponseWrapper {
    private final long startTime;
    private final CountingServletOutputStream output;
    private final PrintWriter writer;

    /**
     * Constructs a response adaptor wrapping the given response.
     *
     * @throws IllegalArgumentException if the response is null
     */
    public CountingServletResponse(HttpServletResponse response) throws IOException {
        super(response);
        startTime = System.currentTimeMillis();
        output = new CountingServletOutputStream(response.getOutputStream());
        writer = new PrintWriter(output, true);
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return output;
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        return writer;
    }

    @Override
    public void flushBuffer() throws IOException {
        writer.flush();
    }

    public long getElapsedTime() {
        return System.currentTimeMillis() - startTime;
    }

    public long getByteCount() throws IOException {
        flushBuffer(); // Ensure that all bytes are written at this point.
        return output.getByteCount();
    }
}
